#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <ctype.h>
#include <getopt.h>

#define OTA_HEADER_SIZE 32
#define OTA_MAGIC1 0xAA
#define OTA_MAGIC2 0x55
#define MAX_TYPE_LEN 10
#define VERSION_DELIMITER '.'

typedef struct {
    uint8_t magic1;          // 固定为 0xAA
    uint8_t magic2;          // 固定为 0x55
    uint8_t firmware_count;  // 固定为 1
    uint8_t major_version;   // 主版本号
    uint8_t minor_version;   // 次版本号
    uint8_t reserved[11];    // 预留字段
    uint32_t firmware_size;  // 固件长度
    uint16_t checksum;       // 校验和
    char firmware_type[10];  // 固件类型
} OtaHeader;

// 函数声明
void print_help();
int pack_file(const char *input_file, const char *type, const char *version);
int unpack_file(const char *input_file);
int view_file_info(const char *input_file);
int convert_hex_to_bin(const char *hex_file, const char *bin_file);
int convert_cly_to_bin(const char *cly_file, const char *bin_file);
uint16_t calculate_checksum(const uint8_t *data, size_t size);
int get_file_extension(const char *filename, char *extension, size_t ext_size);
int create_output_filename(const char *input_file, char *output_file, size_t output_size, const char *suffix, const char *extension);

int main(int argc, char *argv[]) {
    int opt;
    int action = 0; // 0: none, 1: pack, 2: unpack, 3: view

    if (argc < 2) {
        print_help();
        return 1;
    }

    while ((opt = getopt(argc, argv, "hp:u:v:")) != -1) {
        switch (opt) {
            case 'h':
                print_help();
                return 0;
            case 'p':
                if (argc < 5) {
                    printf("Error: Not enough arguments for pack operation\n");
                    print_help();
                    return 1;
                }
                return pack_file(optarg, argv[optind], argv[optind + 1]);
            case 'u':
                return unpack_file(optarg);
            case 'v':
                return view_file_info(optarg);
            default:
                print_help();
                return 1;
        }
    }

    if (action == 0) {
        print_help();
    }

    return 0;
}

void print_help() {
    printf("OTA Pack Tool - A command line tool for OTA firmware packaging\n");
    printf("Version: 1.0\n\n");
    printf("Usage:\n");
    printf("  -h                      Show this help message\n");
    printf("  -p [file] [type] [ver]  Pack firmware file to OTA format\n");
    printf("  -u [file]               Unpack OTA file to binary firmware\n");
    printf("  -v [file]               View OTA file information\n\n");
    printf("Supported input file formats:\n");
    printf("  .bin - Binary file\n");
    printf("  .hex - Intel Hex format file\n");
    printf("  .cly - Custom format file\n\n");
    printf("Examples:\n");
    printf("  otapack -p firmware.bin E10H 1.0\n");
    printf("  otapack -u firmware_ota.bin\n");
    printf("  otapack -v firmware_ota.bin\n");
}

int pack_file(const char *input_file, const char *type, const char *version) {
    FILE *fp_in = NULL, *fp_out = NULL;
    char temp_bin_file[256] = {0};
    char output_file[256] = {0};
    char file_ext[10] = {0};
    uint8_t *firmware_data = NULL;
    OtaHeader header = {0};
    size_t firmware_size = 0;
    int result = 1;
    char *major_str, *minor_str;
    int major, minor;
    
    // 解析版本号
    char version_copy[32];
    strncpy(version_copy, version, sizeof(version_copy) - 1);
    
    major_str = strtok(version_copy, ".");
    minor_str = strtok(NULL, ".");
    
    if (!major_str || !minor_str) {
        printf("Error: Invalid version format. Use format like '1.0'\n");
        return 1;
    }
    
    major = atoi(major_str);
    minor = atoi(minor_str);
    
    if (major < 0 || major > 255 || minor < 0 || minor > 255) {
        printf("Error: Version numbers must be between 0 and 255\n");
        return 1;
    }
    
    // 检查固件类型长度
    if (strlen(type) > MAX_TYPE_LEN) {
        printf("Error: Firmware type too long (max %d characters)\n", MAX_TYPE_LEN);
        return 1;
    }
    
    // 获取文件扩展名
    if (get_file_extension(input_file, file_ext, sizeof(file_ext)) != 0) {
        printf("Error: Could not determine file extension\n");
        return 1;
    }
    
    // 根据文件类型进行处理
    if (strcmp(file_ext, "bin") == 0) {
        // 直接使用bin文件
        strncpy(temp_bin_file, input_file, sizeof(temp_bin_file) - 1);
    } else if (strcmp(file_ext, "hex") == 0) {
        // 转换hex文件为bin文件
        sprintf(temp_bin_file, "%s.temp.bin", input_file);
        if (convert_hex_to_bin(input_file, temp_bin_file) != 0) {
            printf("Error: Failed to convert HEX file to BIN\n");
            return 1;
        }
    } else if (strcmp(file_ext, "cly") == 0) {
        // 转换cly文件为bin文件
        sprintf(temp_bin_file, "%s.temp.bin", input_file);
        if (convert_cly_to_bin(input_file, temp_bin_file) != 0) {
            printf("Error: Failed to convert CLY file to BIN\n");
            return 1;
        }
    } else {
        printf("Error: Unsupported file format: %s\n", file_ext);
        return 1;
    }
    
    // 创建输出文件名
    if (create_output_filename(input_file, output_file, sizeof(output_file), "_ota", "bin") != 0) {
        printf("Error: Failed to create output filename\n");
        goto cleanup;
    }
    
    // 读取固件数据
    fp_in = fopen(temp_bin_file, "rb");
    if (!fp_in) {
        printf("Error: Could not open input file: %s\n", temp_bin_file);
        goto cleanup;
    }
    
    // 获取文件大小
    fseek(fp_in, 0, SEEK_END);
    firmware_size = ftell(fp_in);
    fseek(fp_in, 0, SEEK_SET);
    
    if (firmware_size == 0) {
        printf("Error: Input file is empty\n");
        goto cleanup;
    }
    
    // 分配内存并读取固件数据
    firmware_data = (uint8_t *)malloc(firmware_size);
    if (!firmware_data) {
        printf("Error: Memory allocation failed\n");
        goto cleanup;
    }
    
    if (fread(firmware_data, 1, firmware_size, fp_in) != firmware_size) {
        printf("Error: Failed to read input file\n");
        goto cleanup;
    }
    
    // 计算校验和
    uint16_t checksum = calculate_checksum(firmware_data, firmware_size);
    
    // 准备OTA头
    header.magic1 = OTA_MAGIC1;
    header.magic2 = OTA_MAGIC2;
    header.firmware_count = 1;
    header.major_version = (uint8_t)major;
    header.minor_version = (uint8_t)minor;
    memset(header.reserved, 0, sizeof(header.reserved));
    header.firmware_size = firmware_size;
    header.checksum = checksum;
    strncpy(header.firmware_type, type, MAX_TYPE_LEN);
    
    // 创建输出文件
    fp_out = fopen(output_file, "wb");
    if (!fp_out) {
        printf("Error: Could not create output file: %s\n", output_file);
        goto cleanup;
    }
    
    // 写入OTA头
    if (fwrite(&header, 1, OTA_HEADER_SIZE, fp_out) != OTA_HEADER_SIZE) {
        printf("Error: Failed to write OTA header\n");
        goto cleanup;
    }
    
    // 写入固件数据
    if (fwrite(firmware_data, 1, firmware_size, fp_out) != firmware_size) {
        printf("Error: Failed to write firmware data\n");
        goto cleanup;
    }
    
    printf("Successfully created OTA package: %s\n", output_file);
    printf("Firmware type: %s\n", type);
    printf("Firmware version: %d.%d\n", major, minor);
    printf("Firmware size: %zu bytes\n", firmware_size);
    printf("Checksum: 0x%04X\n", checksum);
    
    result = 0;
    
cleanup:
    if (fp_in) fclose(fp_in);
    if (fp_out) fclose(fp_out);
    if (firmware_data) free(firmware_data);
    
    // 如果创建了临时文件，删除它
    if (strcmp(temp_bin_file, input_file) != 0) {
        remove(temp_bin_file);
    }
    
    return result;
}

int unpack_file(const char *input_file) {
    FILE *fp_in = NULL, *fp_out = NULL;
    char output_file[256] = {0};
    uint8_t *firmware_data = NULL;
    OtaHeader header = {0};
    int result = 1;
    
    // 打开输入文件
    fp_in = fopen(input_file, "rb");
    if (!fp_in) {
        printf("Error: Could not open input file: %s\n", input_file);
        return 1;
    }
    
    // 读取OTA头
    if (fread(&header, 1, OTA_HEADER_SIZE, fp_in) != OTA_HEADER_SIZE) {
        printf("Error: Failed to read OTA header\n");
        goto cleanup;
    }
    
    // 验证OTA头
    if (header.magic1 != OTA_MAGIC1 || header.magic2 != OTA_MAGIC2) {
        printf("Error: Invalid OTA file format\n");
        goto cleanup;
    }
    
    // 创建输出文件名
    if (create_output_filename(input_file, output_file, sizeof(output_file), "_unpacked", "bin") != 0) {
        printf("Error: Failed to create output filename\n");
        goto cleanup;
    }
    
    // 分配内存并读取固件数据
    firmware_data = (uint8_t *)malloc(header.firmware_size);
    if (!firmware_data) {
        printf("Error: Memory allocation failed\n");
        goto cleanup;
    }
    
    if (fread(firmware_data, 1, header.firmware_size, fp_in) != header.firmware_size) {
        printf("Error: Failed to read firmware data\n");
        goto cleanup;
    }
    
    // 验证校验和
    uint16_t calculated_checksum = calculate_checksum(firmware_data, header.firmware_size);
    if (calculated_checksum != header.checksum) {
        printf("Warning: Checksum mismatch. Expected: 0x%04X, Calculated: 0x%04X\n", 
               header.checksum, calculated_checksum);
    }
    
    // 创建输出文件
    fp_out = fopen(output_file, "wb");
    if (!fp_out) {
        printf("Error: Could not create output file: %s\n", output_file);
        goto cleanup;
    }
    
    // 写入固件数据
    if (fwrite(firmware_data, 1, header.firmware_size, fp_out) != header.firmware_size) {
        printf("Error: Failed to write firmware data\n");
        goto cleanup;
    }
    
    printf("Successfully unpacked OTA file: %s\n", output_file);
    printf("Firmware type: %s\n", header.firmware_type);
    printf("Firmware version: %d.%d\n", header.major_version, header.minor_version);
    printf("Firmware size: %u bytes\n", header.firmware_size);
    
    result = 0;
    
cleanup:
    if (fp_in) fclose(fp_in);
    if (fp_out) fclose(fp_out);
    if (firmware_data) free(firmware_data);
    
    return result;
}

int view_file_info(const char *input_file) {
    FILE *fp_in = NULL;
    OtaHeader header = {0};
    
    // 打开输入文件
    fp_in = fopen(input_file, "rb");
    if (!fp_in) {
        printf("Error: Could not open input file: %s\n", input_file);
        return 1;
    }
    
    // 读取OTA头
    if (fread(&header, 1, OTA_HEADER_SIZE, fp_in) != OTA_HEADER_SIZE) {
        printf("Error: Failed to read OTA header\n");
        fclose(fp_in);
        return 1;
    }
    
    // 验证OTA头
    if (header.magic1 != OTA_MAGIC1 || header.magic2 != OTA_MAGIC2) {
        printf("Error: Invalid OTA file format\n");
        fclose(fp_in);
        return 1;
    }
    
    // 显示OTA文件信息
    printf("OTA File Information:\n");
    printf("---------------------\n");
    printf("File: %s\n", input_file);
    printf("Magic: 0x%02X 0x%02X\n", header.magic1, header.magic2);
    printf("Firmware count: %d\n", header.firmware_count);
    printf("Firmware type: %s\n", header.firmware_type);
    printf("Firmware version: %d.%d\n", header.major_version, header.minor_version);
    printf("Firmware size: %u bytes\n", header.firmware_size);
    printf("Checksum: 0x%04X\n", header.checksum);
    
    // 验证校验和
    uint8_t *firmware_data = (uint8_t *)malloc(header.firmware_size);
    if (firmware_data) {
        if (fread(firmware_data, 1, header.firmware_size, fp_in) == header.firmware_size) {
            uint16_t calculated_checksum = calculate_checksum(firmware_data, header.firmware_size);
            printf("Calculated checksum: 0x%04X ", calculated_checksum);
            if (calculated_checksum == header.checksum) {
                printf("(Valid)\n");
            } else {
                printf("(Invalid - checksum mismatch!)\n");
            }
        }
        free(firmware_data);
    }
    
    fclose(fp_in);
    return 0;
}

uint16_t calculate_checksum(const uint8_t *data, size_t size) {
    uint16_t sum = 0;
    for (size_t i = 0; i < size; i++) {
        sum += data[i];
    }
    return sum;
}

int get_file_extension(const char *filename, char *extension, size_t ext_size) {
    const char *dot = strrchr(filename, '.');
    if (!dot || dot == filename) {
        return 1;
    }
    
    strncpy(extension, dot + 1, ext_size - 1);
    extension[ext_size - 1] = '\0';
    
    // 转换为小写
    for (char *p = extension; *p; p++) {
        *p = tolower(*p);
    }
    
    return 0;
}

int create_output_filename(const char *input_file, char *output_file, size_t output_size, 
                          const char *suffix, const char *extension) {
    char base_name[256] = {0};
    const char *dot = strrchr(input_file, '.');
    
    if (!dot || dot == input_file) {
        strncpy(base_name, input_file, sizeof(base_name) - 1);
    } else {
        size_t len = dot - input_file;
        if (len >= sizeof(base_name)) {
            len = sizeof(base_name) - 1;
        }
        strncpy(base_name, input_file, len);
        base_name[len] = '\0';
    }
    
    snprintf(output_file, output_size, "%s%s.%s", base_name, suffix, extension);
    return 0;
}

int convert_hex_to_bin(const char *hex_file, const char *bin_file) {
    char command[512] = {0};
    char temp_bin_file[256] = {0};
    char base_name[256] = {0};
    const char *dot = strrchr(hex_file, '.');
    int result = 1;
    
    // 获取不带后缀的文件名
    if (!dot || dot == hex_file) {
        strncpy(base_name, hex_file, sizeof(base_name) - 1);
    } else {
        size_t len = dot - hex_file;
        if (len >= sizeof(base_name)) {
            len = sizeof(base_name) - 1;
        }
        strncpy(base_name, hex_file, len);
        base_name[len] = '\0';
    }
    
    // hex2bin生成的默认输出文件名
    sprintf(temp_bin_file, "%s.bin", base_name);
    
    // 构建命令行
    sprintf(command, "hex2bin \"%s\"", hex_file);
    
    // 执行命令
    printf("Executing: %s\n", command);
    int cmd_result = system(command);
    
    if (cmd_result != 0) {
        printf("Error: hex2bin command failed with code %d\n", cmd_result);
        return 1;
    }
    
    // 检查输出文件是否存在
    FILE *fp_temp = fopen(temp_bin_file, "rb");
    if (!fp_temp) {
        printf("Error: hex2bin did not create the expected output file: %s\n", temp_bin_file);
        return 1;
    }
    fclose(fp_temp);
    
    // 将hex2bin生成的文件移动到目标位置    
    if (strcmp(temp_bin_file, bin_file) != 0) {
        // 检查目标文件是否存在，如果存在则删除
        FILE *fp_check = fopen(bin_file, "r");
        if (fp_check) {
            fclose(fp_check);
            remove(bin_file);
        }
        
        if (rename(temp_bin_file, bin_file) != 0) {
            printf("Error: Failed to rename hex2bin output file to target location\n");
            return 1;
        }
    }
    
    printf("Successfully converted HEX file to BIN using hex2bin\n");
    return 0;
}

int convert_cly_to_bin(const char *cly_file, const char *bin_file) {
    FILE *fp_in = NULL, *fp_out = NULL;
    char line[1024];
    int line_count = 0;
    uint32_t firmware_size = 0;
    int result = 1;
    
    fp_in = fopen(cly_file, "r");
    if (!fp_in) {
        printf("Error: Could not open CLY file: %s\n", cly_file);
        return 1;
    }
    
    fp_out = fopen(bin_file, "wb");
    if (!fp_out) {
        printf("Error: Could not create output file: %s\n", bin_file);
        fclose(fp_in);
        return 1;
    }
    
    // 读取第一行，获取固件大小
    if (fgets(line, sizeof(line), fp_in) == NULL) {
        printf("Error: Failed to read firmware size from CLY file\n");
        goto cleanup;
    }
    
    // 移除行尾的换行符
    char *newline = strchr(line, '\n');
    if (newline) *newline = '\0';
    
    // 转换为整数
    firmware_size = strtoul(line, NULL, 16);
    
    // 写入固件大小（高字节在前）
    uint8_t size_bytes[4];
    size_bytes[0] = (firmware_size >> 24) & 0xFF;
    size_bytes[1] = (firmware_size >> 16) & 0xFF;
    size_bytes[2] = (firmware_size >> 8) & 0xFF;
    size_bytes[3] = firmware_size & 0xFF;
    
    if (fwrite(size_bytes, 1, 4, fp_out) != 4) {
        printf("Error: Failed to write firmware size to BIN file\n");
        goto cleanup;
    }
    
    // 读取并转换后续行
    line_count = 1;
    while (fgets(line, sizeof(line), fp_in) != NULL) {
        line_count++;
        
        // 移除行尾的换行符
        newline = strchr(line, '\n');
        if (newline) *newline = '\0';
        
        // 检查行长度
        size_t line_len = strlen(line);
        if (line_len % 2 != 0) {
            printf("Error: Invalid line length in CLY file at line %d\n", line_count);
            goto cleanup;
        }
        
        // 转换hex字符串为二进制数据
        uint8_t binary_data[128];
        size_t binary_len = 0;
        
        for (size_t i = 0; i < line_len; i += 2) {
            char hex[3] = {line[i], line[i+1], '\0'};
            binary_data[binary_len++] = (uint8_t)strtoul(hex, NULL, 16);
        }
        
        // 写入二进制数据
        if (fwrite(binary_data, 1, binary_len, fp_out) != binary_len) {
            printf("Error: Failed to write binary data to BIN file at line %d\n", line_count);
            goto cleanup;
        }
    }
    
    printf("Successfully converted CLY file to BIN\n");
    result = 0;
    
cleanup:
    if (fp_in) fclose(fp_in);
    if (fp_out) fclose(fp_out);
    
    return result;
}